package com.tyc.utill;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.UUID;

/**
 * redis + lua 实现限流
 *
 * @author tyc
 * @version 1.0
 * @date 2022-01-22 17:29:58
 */
@Component
public class RedisLimit {
    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    private static final DefaultRedisScript<String> LIMIT_SCRIPT;
    private static final String LIMIT_KEY_PRE = "limit:";

    static {
        LIMIT_SCRIPT = new DefaultRedisScript<>();
        LIMIT_SCRIPT.setLocation(new ClassPathResource("limit.lua"));
        LIMIT_SCRIPT.setResultType(String.class);
    }


    /**
     * uuid作为 value , 当前时间戳作为分数值
     * 每次判断指定时间段内(分值内)value个数
     * 缺点zet的数据越来越大，不推荐使用
     * @return
     */
    public String limitFlow(){
        Long currentTime = System.currentTimeMillis();
        System.out.println(currentTime);
        if(stringRedisTemplate.hasKey("limit")) {
            Integer count = stringRedisTemplate.opsForZSet().rangeByScore("limit", currentTime -  1000*60, currentTime).size();
            System.out.println(count);
            if (count != null && count > 5) {
                return "每分钟最多只能访问5次";
            }
        }
        // Boolean add(K key, V value, double score);
        stringRedisTemplate.opsForZSet().add("limit", UUID.randomUUID().toString(),currentTime);
        return "OK";
    }

    /**
     *  reids + lua 实现限流
     */
    public String limit(String ip,Integer maxTimes){
        // 调用lua脚本
        return stringRedisTemplate.execute(
                LIMIT_SCRIPT,
                Collections.singletonList(LIMIT_KEY_PRE + ip + ":" + System.currentTimeMillis()/1000),
                String.valueOf(maxTimes)
        );
    }


}
